import numpy as np
import json

if __name__ == "__main__":
    # ds_name_l = ['audio']
    ds_name_l = ['audio', 'imagenet', 'movielens', 'music100', 'netflix', 'normal-64', 'text-to-image', 'tiny5m',
                 'word2vec', 'yahoomusic']
    M = 32
    efConstruction = 100
    alpha = 1
    topk = 10
    data_basic_dir = "/home/zhengbian/mips-graph/data"
    graph_basic_dir = "/home/zhengbian/mips-graph/graph"
    res_basic_dir = "/home/zhengbian/mips-graph/result"
    index_cmd_m = {}
    query_cmd_m = {}
    query_no_hier_cmd_m = {}
    for ds_name in ds_name_l:
        base_fname = "%s/%s/%s_base.fvecs" % (data_basic_dir, ds_name, ds_name)
        graph_fname = "%s/%s_%d_%d.hnsw" % (graph_basic_dir, ds_name, M, efConstruction)
        index_cmd = "./ip-nsw/build_index3 %s %s %d %d %.3f" % (base_fname, graph_fname, M, efConstruction, alpha)

        gnd_fname = "%s/%s/%s_top%d_gnd.lshbox" % (data_basic_dir, ds_name, ds_name, topk)
        query_fname = "%s/%s/%s_query.fvecs" % (data_basic_dir, ds_name, ds_name)
        result_fname = "%s/%s.csv" % (res_basic_dir, ds_name)
        query_cmd = "./ip-nsw/query3 %s %s %s %s" % (query_fname, gnd_fname, graph_fname, result_fname)

        gnd_fname = "%s/%s/%s_top%d_gnd.lshbox" % (data_basic_dir, ds_name, ds_name, topk)
        query_fname = "%s/%s/%s_query.fvecs" % (data_basic_dir, ds_name, ds_name)
        result_fname = "%s/%s.csv" % (res_basic_dir, ds_name)
        query_no_hier_cmd = "./ip-nsw/query3_no_hier %s %s %s %s" % (query_fname, gnd_fname, graph_fname, result_fname)

        index_cmd_m[ds_name] = index_cmd
        query_cmd_m[ds_name] = query_cmd
        query_no_hier_cmd_m[ds_name] = query_no_hier_cmd
    with open("index_cmd.json", 'w') as f:
        json.dump(index_cmd_m, f)
    with open("query_cmd.json", 'w') as f:
        json.dump(query_cmd_m, f)
    with open("query_no_hier_cmd.json", 'w') as f:
        json.dump(query_no_hier_cmd_m, f)
